#!/usr/bin/env python3
"""
高级技术指标模块
提供RSI、MACD、KDJ、布林带、ATR等核心技术指标计算
"""

import numpy as np
import pandas as pd
from typing import Dict, List, Tuple, Optional
from dataclasses import dataclass
from enum import Enum


class IndicatorSignal(Enum):
    """技术指标信号"""
    STRONG_BUY = "强烈买入"
    BUY = "买入"
    HOLD = "持有"
    SELL = "卖出"
    STRONG_SELL = "强烈卖出"


@dataclass
class TechnicalIndicatorResult:
    """技术指标结果"""
    name: str
    value: float
    signal: IndicatorSignal
    strength: float  # 信号强度 0-100
    description: str
    support_levels: List[float] = None
    resistance_levels: List[float] = None


class AdvancedTechnicalIndicators:
    """高级技术指标计算器"""
    
    def __init__(self):
        """初始化"""
        pass
    
    def calculate_rsi(self, prices: List[float], period: int = 14) -> TechnicalIndicatorResult:
        """计算RSI相对强弱指标"""
        if len(prices) < period + 1:
            return TechnicalIndicatorResult(
                "RSI", 50.0, IndicatorSignal.HOLD, 50.0, "数据不足"
            )
        
        prices_array = np.array(prices)
        deltas = np.diff(prices_array)
        gains = np.where(deltas > 0, deltas, 0)
        losses = np.where(deltas < 0, -deltas, 0)
        
        # 计算平均收益和损失
        avg_gain = np.mean(gains[-period:])
        avg_loss = np.mean(losses[-period:])
        
        if avg_loss == 0:
            rsi = 100.0
        else:
            rs = avg_gain / avg_loss
            rsi = 100 - (100 / (1 + rs))
        
        # 生成信号
        if rsi >= 80:
            signal, strength = IndicatorSignal.STRONG_SELL, 90.0
        elif rsi >= 70:
            signal, strength = IndicatorSignal.SELL, 70.0
        elif rsi <= 20:
            signal, strength = IndicatorSignal.STRONG_BUY, 90.0
        elif rsi <= 30:
            signal, strength = IndicatorSignal.BUY, 70.0
        else:
            signal, strength = IndicatorSignal.HOLD, 50.0
        
        description = f"RSI={rsi:.1f}, 当前处于{'超买' if rsi > 70 else '超卖' if rsi < 30 else '正常'}区间"
        
        return TechnicalIndicatorResult("RSI", rsi, signal, strength, description)
    
    def calculate_macd(self, prices: List[float], fast: int = 12, slow: int = 26, signal: int = 9) -> TechnicalIndicatorResult:
        """计算MACD指标"""
        if len(prices) < slow + signal:
            return TechnicalIndicatorResult(
                "MACD", 0.0, IndicatorSignal.HOLD, 50.0, "数据不足"
            )
        
        prices_array = np.array(prices)
        
        # 计算EMA
        ema_fast = self._calculate_ema(prices_array, fast)
        ema_slow = self._calculate_ema(prices_array, slow)
        
        # MACD线
        macd_line = ema_fast - ema_slow
        
        # 信号线
        signal_line = self._calculate_ema(macd_line, signal)
        
        # 柱状图
        histogram = macd_line - signal_line
        
        current_macd = macd_line[-1]
        current_signal = signal_line[-1]
        current_hist = histogram[-1]
        
        # 生成信号
        if current_macd > current_signal and current_hist > 0:
            if len(histogram) > 1 and histogram[-2] <= 0:
                signal_type, strength = IndicatorSignal.BUY, 80.0  # 金叉
            else:
                signal_type, strength = IndicatorSignal.BUY, 60.0
        elif current_macd < current_signal and current_hist < 0:
            if len(histogram) > 1 and histogram[-2] >= 0:
                signal_type, strength = IndicatorSignal.SELL, 80.0  # 死叉
            else:
                signal_type, strength = IndicatorSignal.SELL, 60.0
        else:
            signal_type, strength = IndicatorSignal.HOLD, 50.0
        
        description = f"MACD={current_macd:.3f}, 信号线={current_signal:.3f}, 柱状图={'正' if current_hist > 0 else '负'}"
        
        return TechnicalIndicatorResult("MACD", current_macd, signal_type, strength, description)
    
    def calculate_kdj(self, highs: List[float], lows: List[float], closes: List[float], 
                      period: int = 9, k_period: int = 3, d_period: int = 3) -> TechnicalIndicatorResult:
        """计算KDJ指标"""
        if len(closes) < period:
            return TechnicalIndicatorResult(
                "KDJ", 50.0, IndicatorSignal.HOLD, 50.0, "数据不足"
            )
        
        highs_array = np.array(highs)
        lows_array = np.array(lows)
        closes_array = np.array(closes)
        
        # 计算RSV
        rsv_values = []
        for i in range(period - 1, len(closes)):
            highest_high = np.max(highs_array[i - period + 1:i + 1])
            lowest_low = np.min(lows_array[i - period + 1:i + 1])
            
            if highest_high == lowest_low:
                rsv = 50.0
            else:
                rsv = (closes_array[i] - lowest_low) / (highest_high - lowest_low) * 100
            rsv_values.append(rsv)
        
        # 计算K值
        k_values = []
        k = 50.0  # 初始值
        for rsv in rsv_values:
            k = (2 * k + rsv) / 3
            k_values.append(k)
        
        # 计算D值
        d_values = []
        d = 50.0  # 初始值
        for k_val in k_values:
            d = (2 * d + k_val) / 3
            d_values.append(d)
        
        # 计算J值
        j_values = [3 * k - 2 * d for k, d in zip(k_values, d_values)]
        
        current_k = k_values[-1]
        current_d = d_values[-1]
        current_j = j_values[-1]
        
        # 生成信号
        if current_k > current_d and current_k < 20:
            signal_type, strength = IndicatorSignal.STRONG_BUY, 85.0
        elif current_k > current_d and current_k < 50:
            signal_type, strength = IndicatorSignal.BUY, 65.0
        elif current_k < current_d and current_k > 80:
            signal_type, strength = IndicatorSignal.STRONG_SELL, 85.0
        elif current_k < current_d and current_k > 50:
            signal_type, strength = IndicatorSignal.SELL, 65.0
        else:
            signal_type, strength = IndicatorSignal.HOLD, 50.0
        
        description = f"K={current_k:.1f}, D={current_d:.1f}, J={current_j:.1f}"
        
        return TechnicalIndicatorResult("KDJ", current_k, signal_type, strength, description)
    
    def calculate_bollinger_bands(self, prices: List[float], period: int = 20, 
                                  std_dev: float = 2.0) -> TechnicalIndicatorResult:
        """计算布林带指标"""
        if len(prices) < period:
            return TechnicalIndicatorResult(
                "BOLL", 0.5, IndicatorSignal.HOLD, 50.0, "数据不足"
            )
        
        prices_array = np.array(prices)
        
        # 计算中轨（移动平均线）
        middle_band = np.mean(prices_array[-period:])
        
        # 计算标准差
        std = np.std(prices_array[-period:])
        
        # 计算上下轨
        upper_band = middle_band + std_dev * std
        lower_band = middle_band - std_dev * std
        
        current_price = prices_array[-1]
        
        # 计算价格在布林带中的位置
        if upper_band == lower_band:
            position = 0.5
        else:
            position = (current_price - lower_band) / (upper_band - lower_band)
        
        # 生成信号
        if position >= 0.9:
            signal_type, strength = IndicatorSignal.SELL, 80.0  # 接近上轨
        elif position >= 0.75:
            signal_type, strength = IndicatorSignal.SELL, 60.0
        elif position <= 0.1:
            signal_type, strength = IndicatorSignal.BUY, 80.0  # 接近下轨
        elif position <= 0.25:
            signal_type, strength = IndicatorSignal.BUY, 60.0
        else:
            signal_type, strength = IndicatorSignal.HOLD, 50.0
        
        description = f"价格位置={position:.2f}, 上轨={upper_band:.2f}, 下轨={lower_band:.2f}"
        
        return TechnicalIndicatorResult(
            "BOLL", position, signal_type, strength, description,
            support_levels=[lower_band, middle_band],
            resistance_levels=[middle_band, upper_band]
        )
    
    def calculate_atr(self, highs: List[float], lows: List[float], 
                      closes: List[float], period: int = 14) -> TechnicalIndicatorResult:
        """计算ATR真实波动范围"""
        if len(closes) < period + 1:
            return TechnicalIndicatorResult(
                "ATR", 0.0, IndicatorSignal.HOLD, 50.0, "数据不足"
            )
        
        highs_array = np.array(highs)
        lows_array = np.array(lows)
        closes_array = np.array(closes)
        
        # 计算真实范围
        tr_values = []
        for i in range(1, len(closes)):
            tr1 = highs_array[i] - lows_array[i]
            tr2 = abs(highs_array[i] - closes_array[i-1])
            tr3 = abs(lows_array[i] - closes_array[i-1])
            tr = max(tr1, tr2, tr3)
            tr_values.append(tr)
        
        # ATR是真实范围的移动平均
        atr = np.mean(tr_values[-period:])
        current_price = closes_array[-1]
        
        # ATR百分比
        atr_pct = (atr / current_price) * 100 if current_price > 0 else 0
        
        # 波动性评估
        if atr_pct > 5.0:
            volatility = "高波动"
            strength = 80.0
        elif atr_pct > 3.0:
            volatility = "中波动"
            strength = 60.0
        else:
            volatility = "低波动"
            strength = 40.0
        
        description = f"ATR={atr:.3f}, 波动率={atr_pct:.2f}%, {volatility}"
        
        return TechnicalIndicatorResult("ATR", atr_pct, IndicatorSignal.HOLD, strength, description)
    
    def calculate_williams_r(self, highs: List[float], lows: List[float], 
                            closes: List[float], period: int = 14) -> TechnicalIndicatorResult:
        """计算威廉指标"""
        if len(closes) < period:
            return TechnicalIndicatorResult(
                "WR", -50.0, IndicatorSignal.HOLD, 50.0, "数据不足"
            )
        
        highs_array = np.array(highs)
        lows_array = np.array(lows)
        closes_array = np.array(closes)
        
        highest_high = np.max(highs_array[-period:])
        lowest_low = np.min(lows_array[-period:])
        current_close = closes_array[-1]
        
        if highest_high == lowest_low:
            wr = -50.0
        else:
            wr = -100 * (highest_high - current_close) / (highest_high - lowest_low)
        
        # 生成信号
        if wr >= -20:
            signal_type, strength = IndicatorSignal.SELL, 80.0  # 超买
        elif wr >= -30:
            signal_type, strength = IndicatorSignal.SELL, 60.0
        elif wr <= -80:
            signal_type, strength = IndicatorSignal.BUY, 80.0  # 超卖
        elif wr <= -70:
            signal_type, strength = IndicatorSignal.BUY, 60.0
        else:
            signal_type, strength = IndicatorSignal.HOLD, 50.0
        
        description = f"WR={wr:.1f}, 当前处于{'超买' if wr > -20 else '超卖' if wr < -80 else '正常'}区间"
        
        return TechnicalIndicatorResult("WR", wr, signal_type, strength, description)
    
    def calculate_comprehensive_score(self, indicators: Dict[str, TechnicalIndicatorResult]) -> Tuple[float, IndicatorSignal, str]:
        """计算综合技术指标评分"""
        if not indicators:
            return 50.0, IndicatorSignal.HOLD, "无技术指标数据"
        
        # 权重配置
        weights = {
            "RSI": 0.2,
            "MACD": 0.25,
            "KDJ": 0.2,
            "BOLL": 0.15,
            "WR": 0.1,
            "ATR": 0.1
        }
        
        total_score = 0.0
        total_weight = 0.0
        signal_scores = []
        
        for name, indicator in indicators.items():
            weight = weights.get(name, 0.1)
            total_weight += weight
            
            # 将信号转换为分数
            signal_score = self._signal_to_score(indicator.signal)
            weighted_score = signal_score * weight
            total_score += weighted_score
            signal_scores.append(signal_score)
        
        if total_weight > 0:
            final_score = total_score / total_weight
        else:
            final_score = 50.0
        
        # 确定综合信号
        if final_score >= 75:
            final_signal = IndicatorSignal.STRONG_BUY
        elif final_score >= 60:
            final_signal = IndicatorSignal.BUY
        elif final_score <= 25:
            final_signal = IndicatorSignal.STRONG_SELL
        elif final_score <= 40:
            final_signal = IndicatorSignal.SELL
        else:
            final_signal = IndicatorSignal.HOLD
        
        # 生成描述
        active_indicators = len([s for s in signal_scores if s != 50])
        description = f"综合评分{final_score:.1f}, {active_indicators}个指标发出信号"
        
        return final_score, final_signal, description
    
    def _calculate_ema(self, data: np.ndarray, period: int) -> np.ndarray:
        """计算指数移动平均"""
        multiplier = 2 / (period + 1)
        ema = np.zeros_like(data)
        ema[0] = data[0]
        
        for i in range(1, len(data)):
            ema[i] = (data[i] * multiplier) + (ema[i-1] * (1 - multiplier))
        
        return ema
    
    def _signal_to_score(self, signal: IndicatorSignal) -> float:
        """将信号转换为分数"""
        signal_map = {
            IndicatorSignal.STRONG_BUY: 90.0,
            IndicatorSignal.BUY: 70.0,
            IndicatorSignal.HOLD: 50.0,
            IndicatorSignal.SELL: 30.0,
            IndicatorSignal.STRONG_SELL: 10.0
        }
        return signal_map.get(signal, 50.0)


def create_advanced_indicators() -> AdvancedTechnicalIndicators:
    """创建高级技术指标计算器"""
    return AdvancedTechnicalIndicators() 